"""
MQTT协议驱动
支持MQTT消息订阅和发布
基于paho-mqtt库实现
"""

import json
import logging
import threading
import time
from datetime import datetime
from typing import Dict, Any, List, Optional, Callable

from core.base_driver import BaseDriver, TagValue

try:
    import paho.mqtt.client as mqtt
    MQTT_AVAILABLE = True
except ImportError:
    MQTT_AVAILABLE = False

class MqttDriver(BaseDriver):
    """MQTT协议驱动类"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        
        if not MQTT_AVAILABLE:
            raise ImportError("paho-mqtt库未安装，请运行: pip install paho-mqtt")
        
        # MQTT连接参数
        self.host = config.get('host', 'localhost')
        self.port = config.get('port', 1883)
        self.username = config.get('username')
        self.password = config.get('password')
        self.client_id = config.get('client_id', f"iot_collector_{self.device_name}")
        self.keep_alive = config.get('keep_alive', 60)
        self.qos = config.get('qos', 1)
        
        # SSL/TLS配置
        self.use_ssl = config.get('use_ssl', False)
        self.ca_cert = config.get('ca_cert')
        self.cert_file = config.get('cert_file')
        self.key_file = config.get('key_file')
        
        # MQTT客户端
        self.client = None
        
        # 订阅主题和数据缓存
        self.subscribed_topics = {}
        self.data_cache = {}
        self.data_lock = threading.RLock()
        
        # 数据回调
        self.data_callbacks: List[Callable] = []
        
        self.logger.info(f"MQTT驱动初始化完成: {self.host}:{self.port}")
    
    def connect(self) -> bool:
        """建立MQTT连接"""
        try:
            if self.client is None:
                # 使用新的callback API版本
                self.client = mqtt.Client(client_id=self.client_id, callback_api_version=mqtt.CallbackAPIVersion.VERSION2)
                
                # 设置回调函数
                self.client.on_connect = self._on_connect
                self.client.on_disconnect = self._on_disconnect
                self.client.on_message = self._on_message
                self.client.on_subscribe = self._on_subscribe
                self.client.on_publish = self._on_publish
            
            # 设置认证
            if self.username and self.password:
                self.client.username_pw_set(self.username, self.password)
            
            # 设置SSL/TLS
            if self.use_ssl:
                if self.ca_cert:
                    self.client.tls_set(ca_certs=self.ca_cert,
                                      certfile=self.cert_file,
                                      keyfile=self.key_file)
                else:
                    self.client.tls_set()
            
            # 连接MQTT代理
            self.client.connect(self.host, self.port, self.keep_alive)
            
            # 启动网络循环
            self.client.loop_start()
            
            # 等待连接确认
            timeout = 10
            while timeout > 0 and not self.connected:
                time.sleep(0.1)
                timeout -= 0.1
            
            if self.connected:
                self.logger.info(f"MQTT连接成功: {self.host}:{self.port}")
                self._subscribe_to_topics()
                return True
            else:
                self.logger.error("MQTT连接超时")
                return False
                
        except Exception as e:
            self.connected = False
            self.last_error = str(e)
            self.logger.error(f"MQTT连接失败: {e}")
            return False
    
    def disconnect(self):
        """断开MQTT连接"""
        try:
            if self.client:
                self.client.loop_stop()
                self.client.disconnect()
            self.connected = False
            self.logger.info("MQTT连接已断开")
        except Exception as e:
            self.logger.error(f"MQTT断开连接失败: {e}")
    
    def read_tags(self, tags: List[Dict[str, Any]]) -> List[TagValue]:
        """读取标签值（从缓存中获取）"""
        if not self.connected:
            return []
        
        tag_values = []
        
        with self.data_lock:
            for tag in tags:
                tag_name = tag['name']
                topic = tag.get('topic', tag_name)
                
                if topic in self.data_cache:
                    cached_data = self.data_cache[topic]
                    value = cached_data['value']
                    timestamp = cached_data['timestamp']
                    
                    # 应用缩放和转换
                    scaled_value = self.apply_scaling(value, tag.get('scaling', {}))
                    
                    tag_value = TagValue(
                        device_name=self.device_name,
                        tag_name=tag_name,
                        value=scaled_value,
                        timestamp=timestamp,
                        quality="GOOD"
                    )
                    tag_values.append(tag_value)
                else:
                    # 未收到数据
                    tag_value = TagValue(
                        device_name=self.device_name,
                        tag_name=tag_name,
                        value=None,
                        timestamp=datetime.now(),
                        quality="BAD"
                    )
                    tag_values.append(tag_value)
        
        return tag_values
    
    def write_tag(self, tag_name: str, value: Any) -> bool:
        """写入标签值（发布MQTT消息）"""
        if not self.connected:
            return False
        
        # 从配置中找到标签信息
        tag_config = None
        for tag in self.tags:
            if tag['name'] == tag_name:
                tag_config = tag
                break
        
        if not tag_config:
            self.logger.error(f"未找到标签配置: {tag_name}")
            return False
        
        try:
            topic = tag_config.get('publish_topic', tag_config.get('topic', tag_name))
            
            # 准备消息数据
            message_data = self._prepare_message_data(tag_config, value)
            
            # 发布消息
            result = self.client.publish(topic, message_data, qos=self.qos)
            
            if result.rc == mqtt.MQTT_ERR_SUCCESS:
                self.logger.debug(f"MQTT消息发布成功: {topic} = {message_data}")
                return True
            else:
                self.logger.error(f"MQTT消息发布失败: {result.rc}")
                return False
                
        except Exception as e:
            self.logger.error(f"写入标签 {tag_name} 失败: {e}")
            return False
    
    def _subscribe_to_topics(self):
        """订阅MQTT主题"""
        for tag in self.tags:
            topic = tag.get('topic', tag['name'])
            if topic not in self.subscribed_topics:
                try:
                    result = self.client.subscribe(topic, qos=self.qos)
                    if result[0] == mqtt.MQTT_ERR_SUCCESS:
                        self.subscribed_topics[topic] = tag
                        self.logger.debug(f"订阅主题成功: {topic}")
                    else:
                        self.logger.error(f"订阅主题失败: {topic}, 错误码: {result[0]}")
                except Exception as e:
                    self.logger.error(f"订阅主题异常 {topic}: {e}")
    
    def _prepare_message_data(self, tag_config: Dict[str, Any], value: Any) -> str:
        """准备MQTT消息数据"""
        data_format = tag_config.get('data_format', 'json')
        
        if data_format == 'json':
            # JSON格式
            message = {
                'timestamp': datetime.now().isoformat(),
                'device': self.device_name,
                'tag': tag_config['name'],
                'value': value
            }
            return json.dumps(message)
        elif data_format == 'plain':
            # 纯文本格式
            return str(value)
        else:
            # 默认JSON格式
            message = {
                'value': value,
                'timestamp': datetime.now().isoformat()
            }
            return json.dumps(message)
    
    def _parse_message_data(self, topic: str, payload: bytes) -> Any:
        """解析MQTT消息数据"""
        try:
            # 获取对应标签配置
            tag_config = self.subscribed_topics.get(topic)
            if not tag_config:
                return payload.decode('utf-8')
            
            data_format = tag_config.get('data_format', 'json')
            payload_str = payload.decode('utf-8')
            
            if data_format == 'json':
                # JSON格式
                data = json.loads(payload_str)
                
                # 尝试提取value字段
                if isinstance(data, dict):
                    return data.get('value', data)
                else:
                    return data
                    
            elif data_format == 'plain':
                # 纯文本格式
                data_type = tag_config.get('data_type', 'string').lower()
                
                if data_type == 'int':
                    return int(payload_str)
                elif data_type == 'float':
                    return float(payload_str)
                elif data_type == 'bool':
                    return payload_str.lower() in ('true', '1', 'on', 'yes')
                else:
                    return payload_str
            else:
                return payload_str
                
        except Exception as e:
            self.logger.error(f"解析消息失败 {topic}: {e}")
            return payload.decode('utf-8', errors='ignore')
    
    def _on_connect(self, client, userdata, flags, rc, properties=None):
        """连接回调 - 兼容新版API"""
        if rc == 0:
            self.connected = True
            self.last_error = None
            self.logger.info(f"MQTT连接成功: {self.host}:{self.port}")
        else:
            self.connected = False
            self.last_error = f"连接失败，错误代码: {rc}"
            self.logger.error(self.last_error)
    
    def _on_disconnect(self, client, userdata, rc, properties=None):
        """断开连接回调 - 兼容新版API"""
        self.connected = False
        if rc != 0:
            self.logger.warning(f"MQTT意外断开连接: {rc}")
        else:
            self.logger.info("MQTT正常断开连接")
    
    def _on_message(self, client, userdata, msg):
        """消息接收回调"""
        try:
            topic = msg.topic
            payload = msg.payload
            
            self.logger.debug(f"收到MQTT消息: {topic} = {payload}")
            
            # 解析消息数据
            value = self._parse_message_data(topic, payload)
            
            # 缓存数据
            with self.data_lock:
                self.data_cache[topic] = {
                    'value': value,
                    'timestamp': datetime.now(),
                    'qos': msg.qos
                }
            
            # 调用数据回调
            for callback in self.data_callbacks:
                try:
                    callback(topic, value)
                except Exception as e:
                    self.logger.error(f"数据回调执行失败: {e}")
                    
        except Exception as e:
            self.logger.error(f"处理MQTT消息失败: {e}")
    
    def _on_subscribe(self, client, userdata, mid, granted_qos, properties=None):
        """订阅回调 - 兼容新版API"""
        self.logger.debug(f"MQTT订阅成功: mid={mid}, qos={granted_qos}")
    
    def _on_publish(self, client, userdata, mid, properties=None):
        """发布回调 - 兼容新版API"""
        self.logger.debug(f"MQTT消息发布成功: mid={mid}")
    
    def add_data_callback(self, callback: Callable[[str, Any], None]):
        """添加数据回调函数"""
        self.data_callbacks.append(callback)
    
    def remove_data_callback(self, callback: Callable[[str, Any], None]):
        """移除数据回调函数"""
        if callback in self.data_callbacks:
            self.data_callbacks.remove(callback)
    
    def get_device_info(self) -> Dict[str, Any]:
        """获取设备信息"""
        info = super().get_device_info()
        
        with self.data_lock:
            cached_topics = list(self.data_cache.keys())
            
        info.update({
            'host': self.host,
            'port': self.port,
            'client_id': self.client_id,
            'protocol': 'MQTT',
            'use_ssl': self.use_ssl,
            'subscribed_topics': list(self.subscribed_topics.keys()),
            'cached_topics': cached_topics,
            'data_count': len(self.data_cache)
        })
        
        return info
    
    def clear_cache(self):
        """清空数据缓存"""
        with self.data_lock:
            self.data_cache.clear()
        self.logger.info("MQTT数据缓存已清空") 